import matplotlib.pyplot as plt
import numpy as np

def plot_confusion_matrix(confusion_matrix, save_path, class_names=None):
    plt.figure(figsize=(6, 5))
    plt.imshow(confusion_matrix, interpolation='nearest', cmap='Blues')
    plt.colorbar()

    # 类别标签
    if class_names is not None:
        tick_marks = np.arange(len(class_names))
        plt.xticks(tick_marks, class_names)
        plt.yticks(tick_marks, class_names)

    # 在格子中写数值
    thresh = confusion_matrix.max() / 2.
    for i in range(confusion_matrix.shape[0]):
        for j in range(confusion_matrix.shape[1]):
            plt.text(j, i, format(confusion_matrix[i, j], 'd'),
                     ha="center", va="center",
                     color="white" if confusion_matrix[i, j] > thresh else "black")

    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()


def plot_metrics(metrics, save_path):
    plt.figure(figsize=(10, 5))
    plt.plot(metrics['train_loss'], label='Train Loss')
    plt.plot(metrics['val_loss'], label='Val Loss')
    plt.legend()
    plt.savefig(save_path)
    plt.close()

def plot_predicted_vs_true(labels, preds, save_path_prefix):
    # 确保是 numpy 数组
    labels = np.array(labels)
    preds = np.array(preds)

    # ========== 图1: 散点对比图 ==========
    plt.figure(figsize=(6,6))
    plt.scatter(labels, preds, alpha=0.5)
    plt.plot([labels.min(), labels.max()],
             [labels.min(), labels.max()],
             'r--', lw=2)  # 对角线
    plt.xlabel('True Values')
    plt.ylabel('Predicted Values')
    plt.title('Predicted vs True Values')
    plt.axis("equal")
    plt.savefig(f"{save_path_prefix}_scatter.png")
    plt.close()

    # ========== 图2: 分布对比图 ==========
    plt.figure(figsize=(8,6))
    bins = np.linspace(min(labels.min(), preds.min()),
                       max(labels.max(), preds.max()), 50)  # 公共bin
    plt.hist(labels, bins=bins, alpha=0.5, label='True', density=True)
    plt.hist(preds, bins=bins, alpha=0.5, label='Predicted', density=True)
    plt.xlabel("Value")
    plt.ylabel("Density")
    plt.title("Distribution Comparison")
    plt.legend()
    plt.savefig(f"{save_path_prefix}_distribution.png")
    plt.close()


def plot_10fold_loss(all_train_loss, all_val_loss, save_path):
    plt.figure(figsize=(20, 12))  # 超大 figure
    for i in range(10):
        plt.plot(all_train_loss[i], label=f'Train Fold {i+1}', linestyle='-')
        plt.plot(all_val_loss[i], label=f'Val Fold {i+1}', linestyle='--')

    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Train/Val Loss Across 10 Folds')
    plt.legend(ncol=2, fontsize=10)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()

def plot_10fold_confusion_matrix(all_cm, save_path):
    plt.figure(figsize=(20, 12))
    for i, cm in enumerate(all_cm):
        cm = np.array(cm, dtype=int)
        plt.subplot(2, 5, i+1)
        plt.imshow(cm, interpolation='nearest', cmap='Blues')
        plt.title(f'Fold {i+1}')
        plt.colorbar()

        # 在格子里加数值
        thresh = cm.max() / 2.0
        for j in range(cm.shape[0]):
            for k in range(cm.shape[1]):
                plt.text(
                    k, j, f'{cm[j, k]}',
                    ha="center", va="center",
                    color="white" if cm[j, k] > thresh else "black"
                )

    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
    
def plot_10fold_predicted_vs_true(all_labels, all_preds, save_path):
    plt.figure(figsize=(20, 12))
    for i in range(len(all_labels)):
        # 转 numpy 数组，保证有 .min() / .max()
        labels = np.array(all_labels[i])
        preds = np.array(all_preds[i])

        plt.subplot(2, 5, i+1)
        plt.scatter(labels[i], preds[i], alpha=0.5)
        plt.plot([labels[i].min(), labels[i].max()],
                 [labels[i].min(), labels[i].max()],
                 'r--', lw=2)
        plt.xlabel('True Values')
        plt.ylabel('Predicted Values')
        plt.title(f'Fold {i+1}')
        plt.axis("equal")
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()